{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom Estimators\n",
    "\n",
    "In this notebook we'll write an Custom Estimator (using a model function we specifiy). On the way, we'll use ```tf.layers``` to write our model. In the next notebook, we'll use ```tf.layers``` to write a Custom Estimator for a Convolutional Neural Network."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from __future__ import absolute_import\n",
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "\n",
    "import math\n",
    "import numpy as np\n",
    "\n",
    "import tensorflow as tf"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Import the dataset. Here, we'll need to convert the labels to a one-hot encoding, and we'll reshape the MNIST images to (784,)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We'll use Keras (included with TensorFlow) to import the data\n",
    "# I figured I'd do all the preprocessing and reshaping here, \n",
    "# rather than in the model.\n",
    "(x_train, y_train), (x_test, y_test) = tf.contrib.keras.datasets.mnist.load_data()\n",
    "\n",
    "x_train = x_train.astype('float32')\n",
    "x_test = x_test.astype('float32')\n",
    "\n",
    "y_train = y_train.astype('int32')\n",
    "y_test = y_test.astype('int32')\n",
    "\n",
    "# Normalize the color values to 0-1\n",
    "# (as imported, they're 0-255)\n",
    "x_train /= 255\n",
    "x_test /= 255\n",
    "\n",
    "# Flatten 28x28 images to (784,)\n",
    "x_train = x_train.reshape(x_train.shape[0], 784)\n",
    "x_test = x_test.reshape(x_test.shape[0], 784)\n",
    "\n",
    "# Convert to one-hot.\n",
    "y_train = tf.contrib.keras.utils.to_categorical(y_train, num_classes=10)\n",
    "y_test = tf.contrib.keras.utils.to_categorical(y_test, num_classes=10)\n",
    "\n",
    "print(x_train.shape[0], 'train samples')\n",
    "print(x_test.shape[0], 'test samples')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "When using Estimators, we do not manage the TensorFlow session directly. Instead, we skip straight to defining our hyperparameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Number of neurons in each hidden layer\n",
    "HIDDEN1_SIZE = 500\n",
    "HIDDEN2_SIZE = 250"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To write a Custom Estimator we'll specify our own model function. Here, we'll use ```tf.layers``` to replicate the model from the third notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def model_fn(features, labels, mode):\n",
    "    \n",
    "    # First we'll create 2 fully-connected layers, with ReLU activations.\n",
    "    # Notice we're retrieving the 'x' feature (we'll provide this in the input function\n",
    "    # in a moment).\n",
    "    fc1 = tf.layers.dense(features['x'], HIDDEN1_SIZE, activation=tf.nn.relu, name=\"fc1\")\n",
    "    fc2 = tf.layers.dense(fc1, HIDDEN2_SIZE, activation=tf.nn.relu, name=\"fc2\")\n",
    "    \n",
    "    # Add dropout operation; 0.9 probability that a neuron will be kept\n",
    "    dropout = tf.layers.dropout(\n",
    "        inputs=fc2, rate=0.1, training = mode == tf.estimator.ModeKeys.TRAIN, name=\"dropout\")\n",
    "\n",
    "    # Finally, we'll calculate logits. This will be\n",
    "    # the input to our Softmax function. Notice we \n",
    "    # don't apply an activation at this layer.\n",
    "    # If you've commented out the dropout layer,\n",
    "    # switch the input here to 'fc2'.\n",
    "    logits = tf.layers.dense(dropout, units=10, name=\"logits\")\n",
    "    \n",
    "    # Generate Predictions\n",
    "    classes = tf.argmax(logits, axis=1)\n",
    "    predictions = {\n",
    "        'classes': classes,\n",
    "        'probabilities': tf.nn.softmax(logits, name='softmax_tensor')\n",
    "    }\n",
    "    \n",
    "    if mode == tf.estimator.ModeKeys.PREDICT:\n",
    "        # Return an EstimatorSpec for prediction\n",
    "        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)\n",
    "        \n",
    "    # Compute the loss, per usual.\n",
    "    loss = tf.losses.softmax_cross_entropy(\n",
    "        onehot_labels=labels, logits=logits)\n",
    "        \n",
    "    if mode == tf.estimator.ModeKeys.TRAIN:\n",
    "        \n",
    "        # Configure the Training Op\n",
    "        train_op = tf.contrib.layers.optimize_loss(\n",
    "            loss=loss,\n",
    "            global_step=tf.train.get_global_step(),\n",
    "            learning_rate=1e-3,\n",
    "            optimizer='Adam')\n",
    "\n",
    "        # Return an EstimatorSpec for training\n",
    "        return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions,\n",
    "                                      loss=loss, train_op=train_op)    \n",
    "\n",
    "    assert mode == tf.estimator.ModeKeys.EVAL\n",
    "    \n",
    "    # Configure the accuracy metric for evaluation\n",
    "    metrics = {'accuracy': tf.metrics.accuracy(classes, tf.argmax(labels, axis=1))}\n",
    "    \n",
    "    return tf.estimator.EstimatorSpec(mode=mode, \n",
    "                                      predictions=predictions, \n",
    "                                      loss=loss,\n",
    "                                      eval_metric_ops=metrics)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Input functions, as before."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "train_input = tf.estimator.inputs.numpy_input_fn(\n",
    "    {'x': x_train},\n",
    "    y_train, \n",
    "    num_epochs=None, # repeat forever\n",
    "    shuffle=True # \n",
    ")\n",
    "\n",
    "test_input = tf.estimator.inputs.numpy_input_fn(\n",
    "    {'x': x_test},\n",
    "    y_test,\n",
    "    num_epochs=1, # loop through the dataset once\n",
    "    shuffle=False # don't shuffle the test data\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# At this point, our Estimator will work just like a canned one.\n",
    "estimator = tf.estimator.Estimator(model_fn=model_fn)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Train the estimator using our input function.\n",
    "estimator.train(input_fn=train_input, steps=2000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluate the estimator using our input function.\n",
    "# We should see our accuracy metric below\n",
    "evaluation = estimator.evaluate(input_fn=test_input)\n",
    "print(evaluation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MAX_TO_PRINT = 5\n",
    "\n",
    "# This returns a generator object\n",
    "predictions = estimator.predict(input_fn=test_input)\n",
    "i = 0\n",
    "for p in predictions:\n",
    "    true_label = np.argmax(y_test[i])\n",
    "    predicted_label = p['classes']\n",
    "    print(\"Example %d. True: %d, Predicted: %s\" % (i, true_label, predicted_label))\n",
    "    i += 1\n",
    "    if i == MAX_TO_PRINT: break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
